data: {"text": "\u4f60\u597d\uff01\ud83d\ude0a \u6709\u4ec0\u4e48\u6211\u53ef\u4ee5\u5e2e\u4f60\u7684\u5417\uff1f"}

from flask import Flask, request, Response, jsonify
import json
import requests
from collections import defaultdict
import time
import threading
import uuid
import os
import logging

app = Flask(__name__)

# 配置日志
logging.basicConfig(level=logging.INFO, format='[%(asctime)s] %(levelname)s: %(message)s')

# 读取 API_KEY，注意：需要你本地 export API_KEY=xxx
API_KEY = os.getenv("API_KEY", "XIoRyi7beycTNjhQF00003Df68514735B6B78cC08f3d3936")
API_URL = "https://vip.apiyi.com/v1/chat/completions"
MODEL_NAME = "deepseek-v3-250324"

# 配置参数
MAX_HISTORY = 15  # 最多保留15轮对话（实际历史消息数=MAX_HISTORY*2+1）
DEFAULT_MAX_TOKENS = 512
SESSION_TIMEOUT = 1800  

# 会话历史存储器
session_histories = defaultdict(dict)
lock = threading.Lock()

# 定期清理超时会话
def cleanup_expired_sessions():
    while True:
        current_time = time.time()
        expired_sessions = []

        with lock:
            for session_id, data in session_histories.items():
                if current_time - data.get("last_active", 0) > SESSION_TIMEOUT:
                    expired_sessions.append(session_id)

            for session_id in expired_sessions:
                del session_histories[session_id]
                logging.info(f"清理超时会话: {session_id}")

        time.sleep(3600)

# 启动会话清理线程
cleanup_thread = threading.Thread(target=cleanup_expired_sessions, daemon=True)
cleanup_thread.start()

# 初始化日志
logging.info(f"DeepSeek MCP v2.0 | Max History: {MAX_HISTORY} | Timeout: {SESSION_TIMEOUT}s")

@app.route("/mcp", methods=["POST"])
def handle_mcp():
    try:
        data = request.get_json()
        if not data:
            return jsonify({"error": "无效的请求数据"}), 400

        prompt = data.get("prompt", "").strip()
        if not prompt:
            return jsonify({"error": "请输入有效的 prompt"}), 400

        session_id = data.get("session_id")
        max_tokens = data.get("max_tokens", DEFAULT_MAX_TOKENS)
        try:
            max_tokens = int(max_tokens)
            max_tokens = min(max(max_tokens, 16), 2048)
        except ValueError:
            return jsonify({"error": "max_tokens 必须为整数"}), 400

        # 新建会话
        if not session_id:
            session_id = str(uuid.uuid4())
            with lock:
                session_histories[session_id] = {
                    "history": [],
                    "last_active": time.time()
                }

        current_time = time.time()

        # 获取会话历史
        with lock:
            session_data = session_histories.get(session_id)
            if not session_data:
                session_data = {"history": [], "last_active": current_time}
                session_histories[session_id] = session_data

            session_data["last_active"] = current_time
            messages = session_data["history"].copy()

            if not messages:
                messages.append({
                    "role": "system",
                    "content": "You are a helpful assistant that remembers conversation history."
                })

            messages.append({"role": "user", "content": prompt})
            if len(messages) > MAX_HISTORY * 2 + 1:
                messages = messages[-(MAX_HISTORY * 2 + 1):]

        payload = {
            "model": MODEL_NAME,
            "stream": False,
            "max_tokens": max_tokens,
            "messages": messages
        }

        headers = {
            "Content-Type": "application/json",
            "Authorization": f"Bearer {API_KEY}"
        }

        # 发起 API 请求
        resp = requests.post(API_URL, headers=headers, json=payload, timeout=60)
        if resp.status_code != 200:
            logging.error(f"API 调用失败: {resp.status_code} - {resp.text}")
            return jsonify({"error": f"API 调用失败: {resp.status_code}"}), 500

        result = resp.json()
        choices = result.get("choices")
        if not choices or "message" not in choices[0]:
            logging.error(f"API 返回格式异常: {result}")
            return jsonify({"error": "API 返回格式异常"}), 500

        reply = choices[0]["message"]["content"]

        # 更新会话历史
        with lock:
            session_histories[session_id]["history"].append({"role": "user", "content": prompt})
            session_histories[session_id]["history"].append({"role": "assistant", "content": reply})
            session_histories[session_id]["last_active"] = current_time

        service_info = {
            "session_id": session_id,
            "history_count": len(session_histories[session_id]["history"]) // 2,
            "max_tokens": max_tokens
        }

        # 保留 SSE 格式返回
        def sse_format():
            yield f'data: {json.dumps({"text": reply, "service": service_info}, ensure_ascii=False)}\n\n'

        return Response(sse_format(), content_type='text/event-stream')

    except Exception as e:
        logging.exception("服务异常")
        return jsonify({"error": str(e)}), 500

# 健康检查接口
@app.route("/health", methods=["GET"])
def health_check():
    return jsonify({"status": "ok", "active_sessions": len(session_histories)})

if __name__ == "__main__":
    app.run(host="0.0.0.0", port=5001, threaded=True)
